package org.breathe.aof.writer;

import io.netty.buffer.ByteBuf;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

/**
 * @author: breathe
 * @createTime: 2025-05-16
 */
@Slf4j
public class AofBatchWriter {
    //async flush
    private final Writer writer;
    private final ScheduledExecutorService flushScheduler;
    private final AtomicBoolean forceFlush = new AtomicBoolean(false);
    private final int flushInterval;
    //queue
    private static final int DEFAULT_QUEUE_SIZE = 1000;
    private final BlockingQueue<ByteBuf> writeQueue;
    private final Thread writeThread;
    private final AtomicBoolean running = new AtomicBoolean(true);
    //backpressure threshold
    private static final int DEFAULT_BACKPRESSURE_THRESHOLD = 6 * 1024 * 1024;
    private final AtomicLong pendingBytes = new AtomicLong(0);
    public static final int MIN_BATCH_SIZE = 16;
    public static final int MAX_BATCH_SIZE = 50;
    public static final int MIN_BATCH_TIMEOUT_MS = 2;
    public static final int MAX_BATCH_TIMEOUT_MS = 10;
    public static final int LARGE_COMMAND_THRESHOLD = 512 * 1024;
    public AofBatchWriter(Writer writer, int flushInterval) {
        this.writer = writer;
        this.flushInterval = flushInterval;
        this.flushScheduler = new ScheduledThreadPoolExecutor(1, this::createThread);
        this.writeQueue = new LinkedBlockingDeque<>(DEFAULT_QUEUE_SIZE);
        (this.writeThread = createWriteThread()).start();
        if (flushInterval > 0) {
            this.flushScheduler.scheduleAtFixedRate(this::flushTask, flushInterval, flushInterval, TimeUnit.MILLISECONDS);
        }
    }
    private Thread createThread(Runnable r) {
        Thread thread = new Thread(r);
        thread.setName("AOF-Flush-Thread");
        thread.setDaemon(true);
        return thread;
    }
    private Thread createWriteThread() {
        Thread thread = new Thread(this::processWriteQueue);
        thread.setDaemon(true);
        thread.setName("AOF-Writer-Thread");
        return thread;
    }

    private void processWriteQueue() {
        ByteBuf[] batch = new ByteBuf[MAX_BATCH_SIZE];
        int batchSize = 0;
        while (running.get() || !writeQueue.isEmpty()) {
            try {
                int currentBatchSize = calculateBatchSize(batchSize);
                long timeout = calculateTimeout(currentBatchSize);
                long deadline = System.currentTimeMillis() + timeout;
                while (batchSize < currentBatchSize && System.currentTimeMillis() < deadline) {
                    ByteBuf item = writeQueue.poll(Math.max(1, deadline - System.currentTimeMillis()), TimeUnit.MILLISECONDS);
                    if (item != null) {
                        batch[batchSize ++] = item;
                    } else if (batchSize == 0) {
                        Thread.yield();
                    } else {
                        break;
                    }
                }
                if (batchSize > 0) {
                    writeBatch(batch, batchSize);
                    for (int i = 0; i < batchSize; i ++) {
                        pendingBytes.addAndGet(-batch[i].readableBytes());
                        batch[i].release();
                        batch[i] = null;
                    }
                    batchSize = 0;
                }
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                break;
            } catch (Exception e) {
                log.error("fail to process write queue", e);
                for (int i = 0; i < batchSize; i ++) {
                    batch[i].release();
                    batch[i] = null;
                }
                batchSize = 0;
            }
        }
    }

    private void writeBatch(ByteBuf[] batch, int batchSize) {
        if (batchSize <= 0) {
            return ;
        }
        try {
            int totalBytes = 0;
            for (int i = 0; i < batchSize; i ++) {
                totalBytes += batch[i].readableBytes();
            }
            ByteBuffer buffer = ByteBuffer.allocate(totalBytes);
            for (int i = 0; i < batchSize; i ++) {
                buffer.put(batch[i].nioBuffer());
            }
            buffer.flip();
            writer.write(buffer);
        } catch (Exception e) {
            log.error("fail to write batch to AOF file");
        }
    }

    private long calculateTimeout(int currentBatchSize) {
        int queueSize = writeQueue.size();
        return queueSize > DEFAULT_QUEUE_SIZE / 2 ? MIN_BATCH_TIMEOUT_MS : MAX_BATCH_TIMEOUT_MS;
    }

    private int calculateBatchSize(int batchSize) {
        int queueSize = writeQueue.size();
        return Math.min(MAX_BATCH_SIZE, Math.min(MIN_BATCH_SIZE, MIN_BATCH_SIZE + queueSize / 20));
    }

    private void flushTask() {
        try {
            if (forceFlush.compareAndSet(true, false)) {
                writer.flush();
            }
        } catch (Exception e) {
            log.error("fail to flush aof file", e);
        }
    }
    public void write(ByteBuf byteBuf) throws IOException {
        int byteSize = byteBuf.readableBytes();
        if (byteSize > LARGE_COMMAND_THRESHOLD) {
            try {
                ByteBuffer byteBuffer = byteBuf.nioBuffer();
                writer.write(byteBuffer);
            } finally {
                byteBuf.release();
            }
        }
        pendingBytes.addAndGet(byteSize);
        if (pendingBytes.get() > DEFAULT_BACKPRESSURE_THRESHOLD ||
                writeQueue.size() > DEFAULT_QUEUE_SIZE * 0.75) {
            applyBackpressure();
        }
        try {
            boolean success = writeQueue.offer(byteBuf, 3, TimeUnit.SECONDS);
            if (!success) {
                ByteBuffer byteBuffer = byteBuf.nioBuffer();
                writer.write(byteBuffer);
                byteBuf.release();
            }
            forceFlush.set(true);
        } catch (Exception e) {
            byteBuf.release();
            Thread.currentThread().interrupt();
        }
    }

    public void flush() throws IOException {
        while (!writeQueue.isEmpty()) {
            try {
                Thread.sleep(1);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                break;
            }
        }
        writer.flush();
    }

    public void close() throws IOException{
        if (flushScheduler != null) {
            flushScheduler.shutdown();
            try {
                flushScheduler.shutdownNow();
                flushScheduler.awaitTermination(3, TimeUnit.SECONDS);
            } catch (InterruptedException e) {
                flushScheduler.shutdownNow();
                Thread.currentThread().interrupt();
            }
        }
        flush();
        running.set(false);
        writeThread.interrupt();
    }
    private void applyBackpressure() {
        if (pendingBytes.get() > DEFAULT_BACKPRESSURE_THRESHOLD) {
            try {
                long waitTime = Math.min(20, (pendingBytes.get() - DEFAULT_BACKPRESSURE_THRESHOLD) / 1024);
                Thread.sleep(waitTime);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }
        }
    }
}